from langchain_core.chat_history import InMemoryChatMessageHistory
from langchain_core.messages import trim_messages
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableConfig
from langchain_core.runnables.history import RunnableWithMessageHistory

from a0base.base_llm import pop_llm

# 创建聊天提示模板，包含历史消息占位符和用户问题输入
prompt = ChatPromptTemplate.from_messages([
    MessagesPlaceholder("history"),
    ("human", "{question}")
])

# 存储会话历史的字典，键为会话ID，值为对应的历史记录
store = {}

# 设置保留的历史消息轮数（一问一答为一轮）
k = 3

# 配置消息修剪器，用于控制历史消息长度
# 通过消息数量而非实际token数来控制历史长度
trimmer = trim_messages(
    max_tokens=1_000_000,  # 足够大，实际按消息条数控制
    strategy="last",
    # 用消息数量当 token 计数器
    token_counter=lambda msgs: len(msgs),
    include_system=False,
    allow_partial=False,
)


def get_trimmed(session_id):
    """
    获取指定会话ID的修剪后历史消息记录

    Args:
        session_id (str): 会话唯一标识符

    Returns:
        InMemoryChatMessageHistory: 修剪后的历史消息记录对象
    """
    history = store.setdefault(session_id, InMemoryChatMessageHistory())
    history.clear()
    trimmed = trimmer.invoke(history.messages[-k * 2:])
    history.add_messages(trimmed)
    return history


# 创建带消息历史的可运行链对象
chain = RunnableWithMessageHistory(
    prompt | pop_llm,
    get_trimmed,
    input_messages_key="question",
    history_messages_key="history"
)

# 配置运行时参数，设置会话ID
config = RunnableConfig(configurable={"session_id": "demo"})

# 主循环：接收用户输入并输出模型回答
while True:
    question = input("输入问题：")
    response = chain.invoke({"question": question}, config)
    print("回答:", response)
